import numpy as np
import math
import fileModel
from sklearn import metrics


def __PIJ(true_label, pre_label):
    assert len(pre_label) == len(true_label)
    pre_map = {}
    true_map = {}
    for pre_one, true_one in zip(pre_label, true_label):
        if pre_one not in pre_map:
            pre_map[pre_one] = len(pre_map)
        if true_one not in true_map:
            true_map[true_one] = len(true_map)

    # print(len(pre_map), len(true_map))
    matrix = np.zeros((len(pre_map), len(true_map)))
    num_of_pre_label = [0 for item in pre_map]

    for pre_one, true_one in zip(pre_label, true_label):
        num_of_pre_label[pre_map[pre_one]] += 1
        matrix[pre_map[pre_one], true_map[true_one]] += 1

    for i in range(len(pre_map)):
        for j in range(len(true_map)):
            matrix[i, j] /= num_of_pre_label[i]
            matrix[i, j] += 0.00000001

    return matrix, num_of_pre_label


def entropy(true_label, pre_label):
    matrix, cluster_num = __PIJ(true_label, pre_label)
    pre_num, true_num = matrix.shape

    sum = 0
    for i in range(pre_num):
        tmp = 0
        for j in range(true_num):
            # print(matrix[i, j])
            tmp -= matrix[i, j] * math.log2(matrix[i, j])
        # print(tmp)
        sum += tmp * cluster_num[i] / len(pre_label)
    return sum


def purity(true_label, pre_label):
    matrix, cluster_num = __PIJ(true_label, pre_label)
    pre_num, true_num = matrix.shape
    amax = np.amax(matrix, axis=1)

    sum = 0
    for i, value in enumerate(amax):
        # print(value)
        sum += value * cluster_num[i] / len(pre_label)
    # print(i,pre_num)
    assert i+1 == pre_num
    return sum


def get_cluster_from_theta(theta_path, tag_path):
    theta = fileModel.read_float_data(theta_path)
    tag = fileModel.read_one_line_data(tag_path)
    pre_tag = []
    for docIndex,oneDoc in enumerate(theta):
        max = -1
        max_topic_id = -1
        for topic_id,value in enumerate(oneDoc):
            if value > max:
                max_topic_id = topic_id
                max = value
        assert max_topic_id > -1
        pre_tag.append(max_topic_id)

    return tag,pre_tag


def cluster_metric(true_label, pre_label):
    function_list = [metrics.adjusted_rand_score,
                     metrics.v_measure_score,
                     metrics.adjusted_mutual_info_score,
                     metrics.mutual_info_score,
                     metrics.normalized_mutual_info_score,
                     entropy,
                     purity]

    result = []
    for item in function_list:
        result.append(str(item(true_label, pre_label)))

    return result


def evaluation(theta_path, tag_path):
    true_label, pre_label = get_cluster_from_theta(theta_path, tag_path)
    print('\t'.join(cluster_metric( true_label, pre_label)))


if __name__  == '__main__':
    root = "G:/intellij/TopicModelForShortText/My_LDA/data/20newsgroup/"
    for i in range(5):
        kroot = root+"BTM_100/"+str(i)+"/"
        evaluation(kroot+"model-final.theta", root+"20newsgroup.tag")













